import json
import numpy as np
import matplotlib.pyplot as plt
from typing import Dict, List, Tuple
# import matplotlib.pyplot as plt

# 设置全局字体（任选一种）
# plt.rcParams['font.sans-serif'] = ['SimHei']  # 微软雅黑: 'Microsoft YaHei'
# plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题


def calculate_total_traffic(matrix: List[List[float]]) -> float:
    """计算流量矩阵总和（忽略对角线）"""
    np_matrix = np.array(matrix)
    # np.fill_diagonal(np_matrix, 0)  # 确保对角线置零
    return np.sum(np_matrix)


def extract_traffic_totals(data: Dict, pattern: str) -> Tuple[List[float], List[float]]:
    """
    提取指定模式的总流量数据

    参数:
        data: 加载的JSON数据
        pattern: 流量模式名称 (all_to_all等)

    返回:
        (缩放系数列表, 总流量列表) 按缩放系数排序
    """
    entries = data.get(pattern, [])

    # 提取缩放系数和总流量
    scale_factors = []
    totals = []
    for entry in entries:
        # 根据模式键名获取缩放系数
        if "scale_factor" in entry:  # 基础模式
            scale = entry["scale_factor"]
        else:  # 突发组合模式
            scale = entry["base_scale"]

        # 计算总流量
        total = calculate_total_traffic(entry["matrix"])

        scale_factors.append(scale)
        totals.append(total)

    # 按缩放系数排序
    sorted_pairs = sorted(zip(scale_factors, totals), key=lambda x: x[0])
    return zip(*sorted_pairs)


def plot_traffic_totals(data: Dict, output_path: str = "traffic_totals_16.png"):
    """可视化所有模式的总流量趋势"""
    plt.figure(figsize=(10, 6))

    # 定义模式参数
    patterns = {
        # "all_to_all": {
        #     "label": "All-to-All",
        #     "color": "#1f77b4",
        #     # "marker": "o"
        # },
        # "all_to_all_burst": {
        #     "label": "All-to-All + burst",
        #     "color": "#ff7f0e",
        #     "marker": "s"
        # }
        "ring_allreduce": {
            "label": "Ring AllReduce",
            "color": "#800080",
            # "marker": "^"
        },
        "ring_allreduce_burst": {
            "label": "Ring AllReduce + brust",
            "color": "#d62728",
            # "marker": "D"
        }
    }

    # 绘制各模式曲线
    for pattern, style in patterns.items():
        scales, totals = extract_traffic_totals(data, pattern)
        plt.plot(scales, totals,
                 # marker=style["marker"],
                 linestyle='--' if 'burst' in pattern else '-',
                 color=style["color"],
                 linewidth=1,
                 markersize=4,
                 label=style["label"])

    # 图表美化
    plt.xlabel("流量缩放系数", fontsize=12, fontfamily='SimHei')
    plt.ylabel("总流量 (GB)", fontsize=12, fontfamily='SimHei')
    plt.title("不同流量模式的总流量演化(num=16)", fontsize=14, fontfamily='SimHei')
    plt.grid(True, alpha=0.3)
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')

    # 保存输出
    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    plt.show()


# 执行流程
if __name__ == "__main__":
    # 加载数据
    with open("traffic_evolution_16.json", "r") as f:
        traffic_data = json.load(f)

    # 生成可视化
    plot_traffic_totals(traffic_data)